package mr

import (
	"crypto/rand"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"sort"
	"strconv"
	"time"
)
import "log"
import "net/rpc"
import "hash/fnv"

// 生成4位数字的worker ID
func generateWorkerID() string {
	var n uint32
	err := binary.Read(rand.Reader, binary.BigEndian, &n)
	if err != nil {
		panic("Failed to generate worker ID")
	}
	// 取模得到0-9999的随机数
	id := n % 10000
	// 格式化为4位数字字符串，不足补零
	return strconv.FormatUint(uint64(id), 10)
}

var workerID = generateWorkerID()
var workerLog = log.New(os.Stdout, "worker["+workerID+"]: ", log.LstdFlags)

//
// Map functions return a slice of KeyValue.
// 输出的k-v对
//
type KeyValue struct {
	Key   string
	Value string
}

//
// use ihash(key) % NReduce to choose the reduce
// task number for each KeyValue emitted by Map.
//
func ihash(key string) int {
	h := fnv.New32a()                  // 1. 创建FNV-1a 32位哈希算法实例
	h.Write([]byte(key))               // 2. 向哈希器写入key的字节数据
	return int(h.Sum32() & 0x7fffffff) // 3. 计算哈希值并取正整数
}

//请求任务
func getTask() (WorkerTask, error) {
	args := TaskReqArgs{}
	reply := TaskReplyArgs{}

	// 调用RPC服务
	ok := call("Coordinator.RequestTask", &args, &reply)
	if !ok {
		return WorkerTask{}, fmt.Errorf("worker: RPC调用失败: Coordinator.RequestTask") // 不可达代码，仅为编译通过
	}
	if reply.TaskPhase == DonePhase {
		workerLog.Println("所有任务已完成: Coordinator Close")
		return WorkerTask{Phase: DonePhase}, nil
	}

	//初始化worker task
	var task WorkerTask
	task.Phase = reply.TaskPhase
	task.TaskId = reply.TaskId
	task.NMap = reply.NMap
	task.NReduce = reply.NReduce

	//说明是map阶段
	if reply.TaskPhase == MapPhase {
		//没有可用任务，等待下一轮重新请求
		if reply.TaskId == -1 {
			workerLog.Println("当前没有可用任务，等待并重试...")
			time.Sleep(3 * time.Second)
			return getTask()
		}
		task.MapInputFile = reply.FileNames[0]
	} else if reply.TaskPhase == ReducePhase { // reduce阶段
		//没有可用任务，等待下一轮重新请求
		if reply.TaskId == -1 {
			time.Sleep(1 * time.Second)
			return getTask()
		}
		task.ReduceIntermediateFiles = reply.FileNames
		task.OutputFile = reply.OutputFile
	} else {
		return WorkerTask{}, fmt.Errorf("无效任务阶段: %s", reply.TaskPhase)
	}
	return task, nil
}

//执行map任务
func doMap(task WorkerTask, mapf func(string, string) []KeyValue) {
	file := task.MapInputFile
	content, _ := ReadFile(file)
	keyValues := mapf(file, string(content))
	// 自定义格式打印
	workerLog.Printf("生成了 %d 个键值对:", len(keyValues))
	//生成中间文件
	generateIntermediates(task, keyValues)
	workerLog.Printf("map任务%d完成,准备提交", task.TaskId)
}

//执行reduce任务
func doReduce(task WorkerTask, reducef func(string, []string) string) {
	//存放所有的键值对
	var kva = []KeyValue{}
	reduceTaskId := task.TaskId
	nMap := task.NMap

	//获取中间文件
	intermediates := task.ReduceIntermediateFiles
	//读取所有中间文件的数据
	for _, filePath := range intermediates {
		// 打开中间文件
		file, err := os.Open(filePath)
		if err != nil {
			workerLog.Fatalf("无法打开中间文件 %s: %v", file, err)
			return
		}
		defer file.Close() // 确保文件关闭
		//读取中间文件的键值对
		dec := json.NewDecoder(file)
		for {
			var kv KeyValue
			if err := dec.Decode(&kv); err != nil {
				break
			}
			kva = append(kva, kv)
		}
	}

	//之前是在map的单个文件中排序了,但是多个map生成的同一个reduce任务的文件之间没有排序,所以需要重新排序
	sort.Sort(ByKey(kva))

	//生成临时文件
	tmpFile, err := ioutil.TempFile("tmp", fmt.Sprintf("mr-out-%d-tmp-*", reduceTaskId-nMap))
	if err != nil {
		workerLog.Fatal(fmt.Sprintf("创建临时文件失败: %v", err))
	}
	workerLog.Printf("创建临时文件%s成功", tmpFile.Name())
	defer tmpFile.Close()

	//按键相同的分组去调用reducef函数
	i := 0
	for i < len(kva) {
		j := i + 1
		for j < len(kva) && kva[j].Key == kva[i].Key {
			j++
		}
		//把key相同的收集到一个分组
		values := []string{}
		for k := i; k < j; k++ {
			values = append(values, kva[k].Value)
		}
		//调用reducef函数,按格式写入临时文件
		output := reducef(kva[i].Key, values)
		fmt.Fprintf(tmpFile, "%v %v\n", kva[i].Key, output)
		i = j
	}
	// 原子性重命名为标准文件名
	filename := fmt.Sprintf("mr-out-%d", reduceTaskId-nMap)
	if err := os.Rename(tmpFile.Name(), filename); err != nil {
		workerLog.Fatal(fmt.Sprintf("重命名文件失败: %v", err))
	}
	workerLog.Printf("out文件生成完成: %s\n", filename)
	workerLog.Printf("reduce任务%d完成,准备提交", task.TaskId)
}

func submitTask(task WorkerTask) error {
	args := TaskCompleteReqArgs{}
	reply := TaskCompletedReplyArgs{}

	//初始化请求参数
	if task.Phase == MapPhase {
		args.TaskType = MapTask
	} else {
		args.TaskType = ReduceTask
	}

	args.TaskId = task.TaskId

	// 调用RPC服务
	ok := call("Coordinator.CompleteTask", &args, &reply)
	if !ok {
		return fmt.Errorf("worker: RPC调用失败: Coordinator.CompleteTask")
	}

	return nil
}

//定义排序规则
type ByKey []KeyValue

func (a ByKey) Len() int           { return len(a) }
func (a ByKey) Less(i, j int) bool { return a[i].Key < a[j].Key }
func (a ByKey) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }

//生成中间文件
func generateIntermediates(task WorkerTask, keyValues []KeyValue) {
	// 创建临时文件目录（若不存在）
	os.MkdirAll("tmp", 0744)

	mapTaskId := task.TaskId
	nReduce := task.NReduce

	//初始化中间文件的存储数据
	intermediates := make([][]KeyValue, nReduce)
	for reduceId := range intermediates {
		//创建空切片
		intermediates[reduceId] = []KeyValue{}
	}

	//开始根据hash进行分配
	for _, kvs := range keyValues {
		reduceId := ihash(kvs.Key) % nReduce
		intermediates[reduceId] = append(intermediates[reduceId], kvs)
	}

	for reduceTaskId, kvs := range intermediates {
		//排序后再写入
		sort.Sort(ByKey(kvs))
		createJSONIntermediateFile(mapTaskId, reduceTaskId, kvs)
	}
}

//生成临时文件最后原子命名
func createJSONIntermediateFile(mapTaskId int, reduceTaskId int, kvs []KeyValue) {
	//tmpFileName := fmt.Sprintf("mr-%d-%d-tmp-*", mapTaskId, reduceTaskId)
	tmpFile, err := ioutil.TempFile("tmp", fmt.Sprintf("mr-%d-%d-tmp-*", mapTaskId, reduceTaskId))
	if err != nil {
		workerLog.Fatal(fmt.Sprintf("创建临时文件失败: %v", err))
	}
	workerLog.Printf("创建临时文件%s成功", tmpFile.Name())
	defer tmpFile.Close()

	encoder := json.NewEncoder(tmpFile)
	for _, kv := range kvs {
		if err := encoder.Encode(&kv); err != nil {
			workerLog.Fatal(fmt.Sprintf("JSON编码失败: %v", err))
			return
		}
	}
	// 原子性重命名为标准文件名
	filename := fmt.Sprintf("mr-%d-%d", mapTaskId, reduceTaskId)
	if err := os.Rename(tmpFile.Name(), filename); err != nil {
		workerLog.Fatal(fmt.Sprintf("重命名文件失败: %v", err))
	}
	workerLog.Printf("JSON中间文件生成完成: %s，包含 %d 个键值对\n", filename, len(kvs))
}

//
// main/mrworker.go calls this function.
//
func Worker(mapf func(string, string) []KeyValue,
	reducef func(string, []string) string) {

	workerLog.Printf("启动新worker")

	for {
		workerTask, e := getTask()
		if e != nil {
			workerLog.Fatalf("请求出错, %v", e)
			return
		}
		if workerTask.Phase == MapPhase {
			doMap(workerTask, mapf)
			//提交任务，如果出错，直接进入下一个循环，重新请求任务，当前任务由协作者根据超时机制重新分配
			if err := submitTask(workerTask); err != nil {
				workerLog.Fatalf("提交任务出错，%v", err)
			}
		} else if workerTask.Phase == ReducePhase {
			doReduce(workerTask, reducef)
			//提交任务，如果出错，直接进入下一个循环，重新请求任务，当前任务由协作者根据超时机制重新分配
			if err := submitTask(workerTask); err != nil {
				workerLog.Fatalf("提交任务出错，%v", err)
			}
		} else {
			workerLog.Printf("所有任务已完成，协调者已关闭，退出任务")
			break
		}
	}

	// uncomment to send the Example RPC to the coordinator.
	//CallExample()

}

//func CallAssign() (string, string) {
//	args := TaskReqArgs{}
//	reply := TaskReplyArgs{}
//
//	// 调用RPC服务
//	ok := call("Coordinator.ProcessTaskReq", &args, &reply)
//	if !ok {
//		log.Fatalf("RPC调用失败: Coordinator.ProcessTaskReq")
//		return "", "" // 不可达代码，仅为编译通过
//	}
//
//	// 处理任务完成情况
//	if reply.IsDone {
//		log.Printf("全部任务已结束，退出工作流程")
//		return "", ""
//	}
//
//	// 读取文件内容
//	file := reply.FileName
//	data, err := ReadFile(file)
//	if err != nil {
//		log.Fatalf("读取文件失败: %v", err)
//		return "", "" // 不可达代码，仅为编译通过
//	}
//
//	log.Printf("成功获取任务文件: %s, 大小: %d 字节", file, len(data))
//	return file, string(data)
//}

//
// example function to show how to make an RPC call to the coordinator.
//
// the RPC argument and reply types are defined in rpc.go.
//
func CallExample() {
	for i := 1; i <= 10; i++ {
		args := ExampleArgs{}
		args.X = i
		reply := ExampleReply{}
		ok := call("Coordinator.Example", &args, &reply)
		if ok {
			// reply.Y should be 100.
			fmt.Printf("reply.Y %v\n", reply.Y)
		} else {
			fmt.Printf("call failed!\n")
		}
	}
	args := ExampleArgs{X: -1}
	reply := ExampleReply{}
	call("Coordinator.Example", &args, &reply)
	//// declare an argument structure.
	//args := ExampleArgs{}
	//
	//// fill in the argument(s).
	//args.X = 99
	//
	//// declare a reply structure.
	//reply := ExampleReply{}
	//
	//// send the RPC request, wait for the reply.
	//// the "Coordinator.Example" tells the
	//// receiving server that we'd like to call
	//// the Example() method of struct Coordinator.
	//ok := call("Coordinator.Example", &args, &reply)
	//if ok {
	//	// reply.Y should be 100.
	//	fmt.Printf("reply.Y %v\n", reply.Y)
	//} else {
	//	fmt.Printf("call failed!\n")
	//}
}

//
// send an RPC request to the coordinator, wait for the response.
// usually returns true.
// returns false if something goes wrong.
//
func call(rpcname string, args interface{}, reply interface{}) bool {
	// c, err := rpc.DialHTTP("tcp", "127.0.0.1"+":1234")
	sockname := coordinatorSock()
	c, err := rpc.DialHTTP("unix", sockname)
	if err != nil {
		workerLog.Fatal("dialing:", err)
	}
	defer c.Close()

	//c：代表一个已连接到远程 RPC 服务器的客户端
	// 通过这个客户端，可以调用远程服务器上的方法
	err = c.Call(rpcname, args, reply)
	if err == nil {
		return true
	}

	fmt.Println(err)
	return false
}
